package com.zengde.ai_sandbox.controller;

import com.zengde.ai_sandbox.service.QdrantService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.*;

@RestController
@RequestMapping("/api/qdrant")
@RequiredArgsConstructor
@Slf4j
public class QdrantDemoController {

    private final QdrantService qdrantService;

    /**
     * 初始化Qdrant演示
     */
    @PostMapping("/init")
    public ResponseEntity<Map<String, String>> initQdrant() {
        try {
            log.info("开始初始化Qdrant演示");
            
            // 检查服务是否可用
            if (!qdrantService.isQdrantAvailable()) {
                Map<String, String> response = new HashMap<>();
                response.put("status", "error");
                response.put("message", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
                return ResponseEntity.internalServerError().body(response);
            }
            
            // 创建集合
            qdrantService.createCollection();
            
            // 添加示例向量
            qdrantService.addVectors();
            
            Map<String, String> response = new HashMap<>();
            response.put("status", "success");
            response.put("message", "Qdrant演示初始化成功，已创建集合并添加示例数据");
            
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            log.error("初始化Qdrant演示时出错: {}", e.getMessage(), e);
            
            Map<String, String> response = new HashMap<>();
            response.put("status", "error");
            response.put("message", "初始化失败: " + e.getMessage());
            
            return ResponseEntity.internalServerError().body(response);
        }
    }

    /**
     * 搜索相似向量
     */
    @GetMapping("/search")
    public ResponseEntity<Map<String, Object>> searchSimilar(
            @RequestParam(defaultValue = "3") int limit) {
        try {
            log.info("开始搜索相似向量，限制结果数量: {}", limit);
            
            // 检查服务是否可用
            if (!qdrantService.isQdrantAvailable()) {
                Map<String, Object> response = new HashMap<>();
                response.put("status", "error");
                response.put("message", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
                return ResponseEntity.internalServerError().body(response);
            }
            
            // 使用示例查询向量
            List<Float> queryVector = List.of(0.2f, 0.3f, 0.4f, 0.5f);
            
            // 执行搜索
            List<Map<String, Object>> results = qdrantService.searchSimilar(queryVector, limit);
            
            Map<String, Object> response = new HashMap<>();
            response.put("status", "success");
            response.put("query_vector", queryVector);
            response.put("results", results);
            
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            log.error("搜索相似向量时出错: {}", e.getMessage(), e);
            
            Map<String, Object> response = new HashMap<>();
            response.put("status", "error");
            response.put("message", "搜索失败: " + e.getMessage());
            
            return ResponseEntity.internalServerError().body(response);
        }
    }

    /**
     * 健康检查接口
     */
    @GetMapping("/health")
    public ResponseEntity<Map<String, String>> healthCheck() {
        Map<String, String> response = new HashMap<>();
        
        if (qdrantService.isQdrantAvailable()) {
            response.put("status", "up");
            response.put("qdrant_service", "available");
        } else {
            response.put("status", "degraded");
            response.put("qdrant_service", "unavailable");
            response.put("message", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        response.put("service", "qdrant-demo");
        response.put("timestamp", new Date().toString());
        
        return ResponseEntity.ok(response);
    }
} 